from langchain_community.llms import Tongyi
from langchain.agents import Tool,create_react_agent,AgentExecutor,tool
from langchain.prompts import ChatPromptTemplate
from .models import Room,Reservation




llm = Tongyi()



@tool
def fun1(input):
    """
    当用户想要查询酒店的房间时使用这个函数
    """
    data = list(Room.objects.all().values())
    return data

def fun2(input):
    data = list(Reservation.objects.all().values())
    return data

def fun3(input):
    return "七天无理由退货"


tools = [
    Tool(func=fun2,
         name='fun2',
         description='用户询问订单问题，使用这个函数'),
    Tool(func=fun3,
         name='fun3',
         description='用户咨询售后，使用这个函数'),
    fun1
]


template = '''Answer the following questions as best you can. You have access to the following tools:

            {tools}

            Use the following format:
            Question: the input question you must answer
            Thought: you should always think about what to do
            Action: the action to take, should be one of [{tool_names}]
            Action Input: the input to the action
            Observation: the result of the action
            ... (this Thought/Action/Action Input/Observation can repeat N times)
            Final Answer: the answer to the final question asked

            Begin!

            Question: {input}
            Thought:{agent_scratchpad}'''


prompt = ChatPromptTemplate.from_template(template)

agent = create_react_agent(llm=llm, tools=tools, prompt=prompt)
agent_executor = AgentExecutor(agent=agent, tools=tools,verbose=True)

def generate_sse_lc(chunks):
    print("####")
    for chunk in chunks:
        if chunk:
            data = f"data: {chunk}\n\n"
            yield data.encode('utf-8')
        else:
            print('____________')
            return 'no message'

def method(input):
    return agent_executor.invoke({"input": input})
    # chunks = agent_executor.stream({"input": input})
    # response = StreamingHttpResponse(
    #     generate_sse_lc(chunks),
    #     content_type="text/event-stream",
    # )
    # response["Cache-Control"] = "no-cache"
    # return response